import lightning as L
from torch.utils.data import DataLoader
from omegaconf import DictConfig

from ecgcmr.imaging.img_dataset.ContrastiveImagingDataset import ContrastiveImageDataset
from ecgcmr.imaging.img_dataset.DownstreamImageDataset import DownstreamImageDataset


class ContrastiveImageDataModule(L.LightningDataModule):
    def __init__(
        self,
        cfg: DictConfig
        ) -> None:
        super().__init__()

        self.cfg = cfg
        self.batch_size = cfg.dataset.batch_size
        self.num_workers = cfg.dataset.num_workers

    def setup(self, stage: str):
        if stage == 'fit':
            self.dataset_train = ContrastiveImageDataset(cfg=self.cfg, mode='train', apply_augmentations=True)
            self.dataset_val = ContrastiveImageDataset(cfg=self.cfg, mode='val', apply_augmentations=False)

    def train_dataloader(self):
        return DataLoader(
            self.dataset_train, 
            batch_size=self.batch_size, 
            shuffle=True,
            num_workers=self.num_workers,
            pin_memory=True
        )

    def val_dataloader(self):
        return DataLoader(
            self.dataset_val, 
            batch_size=self.batch_size, 
            shuffle=False,
            num_workers=self.num_workers,
            pin_memory=True
        )

class ContrastiveImageDataModule_PyTorch:
    def __init__(
        self,
        cfg: DictConfig
        ) -> None:
        super().__init__()

        self.cfg = cfg
        self.main_batch_size = cfg.dataset.batch_size
        self.downstream_batch_size = cfg.downstream_task.batch_size
        self.num_workers = cfg.dataset.num_workers

    def setup(self, stage: str):
        if stage == 'fit':
            self.dataset_train_main = ContrastiveImageDataset(cfg=self.cfg, mode='train', apply_augmentations=True)
            self.dataset_train_downstream = DownstreamImageDataset(cfg=self.cfg, mode='train', apply_augmentations=False)

            self.dataset_val_main = ContrastiveImageDataset(cfg=self.cfg, mode='val', apply_augmentations=True)
            self.dataset_val_downstream = DownstreamImageDataset(cfg=self.cfg, mode='val', apply_augmentations=False)

    def train_dataloader(self):
        loader_main = DataLoader(
            self.dataset_train_main, 
            batch_size=self.main_batch_size, 
            shuffle=True, 
            num_workers=self.num_workers,
            pin_memory=True
        )

        loader_downstream = DataLoader(
            self.dataset_train_downstream,
            batch_size=self.downstream_batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            pin_memory=True
        )

        return {'main': loader_main, 'downstream': loader_downstream}

    def val_dataloader(self):
        loader_downstream = DataLoader(
            self.dataset_val_downstream,
            batch_size=self.main_batch_size,
            shuffle=False,
            num_workers=self.num_workers,
            pin_memory=True
        )

        loader_main = DataLoader(
            self.dataset_val_main, 
            batch_size=self.main_batch_size, 
            shuffle=False,
            num_workers=self.num_workers,
            pin_memory=True
        )

        return {'main': loader_main, 'downstream': loader_downstream}
